import glob
import os, sys
import time
import numpy as np
from astropy.io import fits
from astropy.table import Table
from astropy.time import Time
from matplotlib import pyplot as plt
from scipy.ndimage import median_filter
from tqdm import trange


def make_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)
    return path
    
def create_folder(path):
    if not os.path.exists(path):
        os.makedirs(path)
    return path

def clean_dir(path, suffix='.fits'):
    flist = glob.glob('{}/*{}'.format(path, suffix))
    for f in flist:
        os.remove(f)


def remove_dir(path):
    clean_dir(path)
    os.removedirs(path)


def remove_file(path):
    if os.path.exists(path):
        os.remove(path)


def timetag(txt, color=44):
    t = time.asctime(time.localtime(time.time()))
    print('\033[37;{}m{}: {}\033[0m'.format(color, t, txt))


def hltxt(txt, color=42):
    print('\033[37;{}m{}\033[0m'.format(color, txt))


def get_block_limits(n, nblock):
    ncols = n // nblock
    cols = np.zeros(nblock, dtype=int) + ncols
    res = n - ncols * nblock
    if res != 0:
        cols[0:res] = cols[0:res] + 1
    limits = np.cumsum(cols)
    limits = np.append([0], limits)
    return limits


def read_header(header, key, default_value):
    if key in header:
        return header[key]
    return default_value


def make_image_cat(flist, root=None, stats=True):
    namelist = []
    datetag, obstype, tstart, tsys, wave = [], [], [], [], []
    texp, tdark, imgid, seqnum, nstack = [], [], [], [], []
    dn_mean, dn_std = [], []
    for i in trange(len(flist), desc='indexing images'):
        f = flist[i]
        # file name and root directory
        if root is None:
            namelist.append(f)
        else:
            if not f.startswith(root):
                raise Exception('Invalid root directory: {}'.format(root))
            tt = f[len(root):]
            if tt.startswith('/'):
                tt = tt[1:]
            namelist.append(tt)

        # headers
        with fits.open(f) as img:
            header = img[0].header
            datetag.append(read_header(header, 'DATETAG', 'NULL'))
            obstype.append(read_header(header, 'OBSTYPE', 'NULL'))
            tsys.append(read_header(header, 'SYSTEMP', 0.0))
            wave.append(read_header(header, 'WAVELEN', 0.0))
            tstart.append(read_header(header, 'DATE-OBS', 'NULL'))
            texp.append(read_header(header, 'EXPTIME', -99.0))
            tdark.append(read_header(header, 'DARKTIME', -99.0))
            imgid.append(read_header(header, 'IMGINDEX', 0))
            seqnum.append(read_header(header, 'SEQNUM', 0))
            nstack.append(read_header(header, 'NSTACK', 0))
            if stats:
                dn_mean.append(np.mean(img[0].data))
                dn_std.append(np.std(img[0].data))

    tab = Table()
    tab['name'] = namelist
    tab['tag'] = datetag
    tab['obstype'] = obstype
    tab['t'] = tstart
    tab['tsys'] = np.array(tsys).astype(np.float32)
    tab['wave'] = np.array(wave).astype(np.float32)
    tab['texp'] = np.array(texp).astype(np.float32)
    tab['tdark'] = np.array(tdark).astype(np.float32)
    tab['imgid'] = np.array(imgid).astype(int)
    tab['seqnum'] = np.array(seqnum).astype(int)
    tab['nstack'] = np.array(nstack).astype(int)
    if stats:
        tab['mean'] = np.array(dn_mean).astype(np.float32)
        tab['std'] = np.array(dn_std).astype(np.float32)
    tab.sort(['tsys', 'wave', 'tag', 'texp', 'tdark', 'imgid', 'seqnum', 'name'])
    if root is None:
        tab.meta['keywords'] = {'root': {'value': ''}}
    else:
        tab.meta['keywords'] = {'root': {'value': root}}
    return tab


def get_flist(cat):
    flist = np.array(cat['name'])
    if 'keywords' in cat.meta:
        if 'root' in cat.meta['keywords']:
            if 'value' in cat.meta['keywords']['root']:
                root = cat.meta['keywords']['root']['value']
                if root != '':
                    flist = np.array(['{}/{}'.format(root, f) for f in flist])
    return flist


def row_to_header(row):
    cols = row.colnames
    header = dict()
    for c in cols:
        if c == 'tag':
            header['DATETAG'] = row[c]
        if c == 'obstype':
            header['OBSTYPE'] = row[c]
        if c == 'tclass':
            header['SYSTEMP'] = row[c]
        if c == 'wave':
            header['WAVELEN'] = row[c]
        if c == 'imgid':
            header['IMGINDEX'] = row[c]
        if c == 'seqnum':
            header['SEQNUM'] = row[c]
    return header


def get_tag(value, default_value):
    if value == 'NULL':
        return default_value
    return value.replace('-', '_')


def class_tsys(tsys, t0):
    """
        根据工况温度T0生成一个温度列表，默认温差间隔是5K
        @2023-05-17:临时将温度间隔改为3K，并将温度列表的长度更改为10 (by XYH)
    """
    dT = 5   # default value: 5
    nT = 10  # default value 10
    if 'DELTA_TEMP' in os.environ:
        dT = int(os.environ['DELTA_TEMP'])
    ts = (np.arange(nT, dtype=int) - 3) * dT + t0
    tclass = []
    for t in tsys:
        index = np.argmin(np.abs(ts - t))
        tclass.append((index - 3) * dT + t0)
#    print('TCLASS: {}'.format(tclass))
    return tclass


def match_dark(cat, tdark):
    index = cat['tdark'] <= tdark
    if np.sum(index) == 0:
        index = np.argmin(cat['tdark'])
        return get_flist(cat)[index]
    else:
        ind = np.argmax(cat['tdark'][index])
        return get_flist(cat)[index][ind]


def autoflip_monitor(data):
    index = np.argmax(np.abs(data))
    if data[index] < 0:
        return data * (-1)
    else:
        return data


def read_monitor(ftxt, smooth=2, autoflip=False, ffig=None, montype='sphere'):
    cat = Table.read(ftxt, format='ascii.no_header')
    cat['t'] = cat['col1']
    cat['mjd'] = Time(cat['col1']).mjd
    cat['mjd'].unit = 'day'
    cat['raw'] = cat['col2'].copy()
    cat['raw'].unit = 'A'
    cat['value'] = cat['col2'].copy()
    cat = cat['t', 'mjd', 'raw', 'value']
    cat.sort('mjd')
    if autoflip:
        cat['value'] = autoflip_monitor(cat['value'])
    else:
        cat['value'] = cat['value'] * (-1)
    not_smoothed = cat['value'].copy()
    if smooth > 0:
        cat['value'] = median_filter(cat['value'], smooth * 2 + 1)
    cat['value'].unit = 'A'
    cat.meta['keywords'] = {'type': {'value': montype}}

    if ffig is not None:
        plt.figure(figsize=(12, 5))
        plt.subplot(121)
        plt.scatter(cat['mjd'], cat['raw'], s=1, color='grey')
        plt.xlabel('modified Julian date')
        plt.ylabel('original reading (A)')

        plt.subplot(122)
        t = (cat['mjd'] - cat['mjd'][0]) * 24
        if smooth > 0:
            plt.scatter(t, not_smoothed, s=1, color='grey')
        plt.scatter(t, cat['value'], s=1, color='g')
        plt.xlabel('hours from start')
        plt.ylabel('value (A)')

        plt.tight_layout()
        plt.savefig(ffig)
        plt.close()

    return cat


def get_monitor_value(t, value, tstart, tend, method=np.median):
    index = (t > tstart) * (t < tend)
    if np.sum(index) == 0:
        result = np.interp((tstart + tend) / 2, t, value, left=np.nan, right=np.inf)
    else:
        result = method(value[index])
    return result

#add on 20250311
def get_current(mon_tab,tmjds,fmon):
    ts = mon_tab['mjd']
    vals = mon_tab['value']
    currents = []
    plt.scatter(ts,vals,s=0.5)
    for tmjd in tmjds:
        plt.axvline(x=tmjd,linestyle='--',color='green')
        #plt.show()
        idx = ts>tmjd
        idx *= vals>1e-9
        idxmin = np.min(np.where(idx)[0])
        print("idxmin",idxmin)
        idxmax = idxmin
        while(vals[idxmax]>1e-9):
            idxmax += 1
            if idxmax == len(vals)-1:
                break
        idxmax -= 1
        plt.axvline(x=ts[idxmin],linestyle='--',color='red')
        plt.axvline(x=ts[idxmax],linestyle='--',color='black')
        currents.append(np.mean(vals[idxmin:idxmax+1]))
    plt.savefig(fmon,bbox_inches='tight')
    plt.show()
    plt.clf()
    return currents
        
def read_switchtime(fname):
    cat = Table.read(fname, format='ascii.no_header')
    cat['t'] = cat['col1']
    cat['mjd'] = Time(cat['col1']).mjd
    cat['mjd'].unit = 'day'
    cat['wave'] = cat['col2'].astype(np.float32)
    cat['wave'].unit = 'nm'
    cat = cat['t', 'mjd', 'wave']
    return cat
